// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <nebula/array/builder_binary.h>

#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <cstring>
#include <numeric>
#include <string>
#include <utility>
#include <vector>

#include <nebula/core/array.h>
#include <nebula/core/buffer.h>
#include <turbo/utility/status.h>
#include <nebula/types/type.h>
#include <nebula/types/type_traits.h>
#include <nebula/bits/bit_util.h>
#include <turbo/base/checked_cast.h>
#include <nebula/numeric/decimal.h>
#include <turbo/log/logging.h>
#include <nebula/core/visit_data_inline.h>

namespace nebula {


    // ----------------------------------------------------------------------
    // Binary/StringView
    BinaryViewBuilder::BinaryViewBuilder(const std::shared_ptr<DataType> &type,
                                         MemoryPool *pool)
            : BinaryViewBuilder(pool) {}

    turbo::Status BinaryViewBuilder::append_array_slice(const ArraySpan &array, int64_t offset,
                                                      int64_t length) {
        auto bitmap = array.get_values<uint8_t>(0, 0);
        auto values = array.get_values<BinaryViewType::c_type>(1) + offset;

        int64_t out_of_line_total = 0, i = 0;
        VisitNullBitmapInline(
                array.buffers[0].data, array.offset + offset, length, array.null_count,
                [&] {
                    if (!values[i].is_inline()) {
                        out_of_line_total += static_cast<int64_t>(values[i].size());
                    }
                    ++i;
                },
                [&] { ++i; });

        TURBO_RETURN_NOT_OK(Reserve(length));
        TURBO_RETURN_NOT_OK(reserve_data(out_of_line_total));

        for (int64_t i = 0; i < length; i++) {
            if (bitmap && !bit_util::get_bit(bitmap, array.offset + offset + i)) {
                unsafe_append_null();
                continue;
            }

            unsafe_append(util::FromBinaryView(values[i], array.get_variadic_buffers().data()));
        }
        return turbo::OkStatus();
    }

    turbo::Status BinaryViewBuilder::finish_internal(std::shared_ptr<ArrayData> *out) {
        TURBO_MOVE_OR_RAISE(auto null_bitmap, null_bitmap_builder_.FinishWithLength(length_));
        TURBO_MOVE_OR_RAISE(auto data, data_builder_.FinishWithLength(length_));
        TURBO_MOVE_OR_RAISE(auto byte_buffers, data_heap_builder_.finish());
        BufferVector buffers(byte_buffers.size() + 2);
        buffers[0] = std::move(null_bitmap);
        buffers[1] = std::move(data);
        std::move(byte_buffers.begin(), byte_buffers.end(), buffers.begin() + 2);
        *out = ArrayData::create(type(), length_, std::move(buffers), null_count_);
        reset();
        return turbo::OkStatus();
    }

    turbo::Status BinaryViewBuilder::reserve_data(int64_t length) {
        return data_heap_builder_.Reserve(length);
    }

    void BinaryViewBuilder::reset() {
        ArrayBuilder::reset();
        data_builder_.reset();
        data_heap_builder_.reset();
    }

    // ----------------------------------------------------------------------
    // Fixed width binary

    FixedSizeBinaryBuilder::FixedSizeBinaryBuilder(const std::shared_ptr<DataType> &type,
                                                   MemoryPool *pool, int64_t alignment)
            : ArrayBuilder(pool, alignment),
              byte_width_(turbo::checked_cast<const FixedSizeBinaryType &>(*type).byte_width()),
              byte_builder_(pool, alignment) {}

    void FixedSizeBinaryBuilder::check_value_size(int64_t size) {
                DKCHECK_EQ(size, byte_width_) << "Appending wrong size to FixedSizeBinaryBuilder";
    }

    turbo::Status FixedSizeBinaryBuilder::append_values(const uint8_t *data, int64_t length,
                                                       const uint8_t *valid_bytes) {
        TURBO_RETURN_NOT_OK(Reserve(length));
        unsafe_append_to_bitmap(valid_bytes, length);
        return byte_builder_.append(data, length * byte_width_);
    }

    turbo::Status FixedSizeBinaryBuilder::append_values(const uint8_t *data, int64_t length,
                                                       const uint8_t *validity,
                                                       int64_t bitmap_offset) {
        TURBO_RETURN_NOT_OK(Reserve(length));
        unsafe_append_to_bitmap(validity, bitmap_offset, length);
        return byte_builder_.append(data, length * byte_width_);
    }

    turbo::Status FixedSizeBinaryBuilder::append_null() {
        TURBO_RETURN_NOT_OK(Reserve(1));
        unsafe_append_null();
        return turbo::OkStatus();
    }

    turbo::Status FixedSizeBinaryBuilder::append_nulls(int64_t length) {
        TURBO_RETURN_NOT_OK(Reserve(length));
        unsafe_append_to_bitmap(length, false);
        byte_builder_.unsafe_append(/*num_copies=*/length * byte_width_, 0);
        return turbo::OkStatus();
    }

    turbo::Status FixedSizeBinaryBuilder::append_empty_value() {
        TURBO_RETURN_NOT_OK(Reserve(1));
        unsafe_append_to_bitmap(true);
        byte_builder_.unsafe_append(/*num_copies=*/byte_width_, 0);
        return turbo::OkStatus();
    }

    turbo::Status FixedSizeBinaryBuilder::append_empty_values(int64_t length) {
        TURBO_RETURN_NOT_OK(Reserve(length));
        unsafe_append_to_bitmap(length, true);
        byte_builder_.unsafe_append(/*num_copies=*/length * byte_width_, 0);
        return turbo::OkStatus();
    }

    void FixedSizeBinaryBuilder::reset() {
        ArrayBuilder::reset();
        byte_builder_.reset();
    }

    turbo::Status FixedSizeBinaryBuilder::resize(int64_t capacity) {
        TURBO_RETURN_NOT_OK(check_capacity(capacity));
        TURBO_RETURN_NOT_OK(byte_builder_.resize(capacity * byte_width_));
        return ArrayBuilder::resize(capacity);
    }

    turbo::Status FixedSizeBinaryBuilder::finish_internal(std::shared_ptr<ArrayData> *out) {
        std::shared_ptr<Buffer> data;
        TURBO_RETURN_NOT_OK(byte_builder_.finish(&data));

        std::shared_ptr<Buffer> null_bitmap;
        TURBO_RETURN_NOT_OK(null_bitmap_builder_.finish(&null_bitmap));
        *out = ArrayData::create(type(), length_, {null_bitmap, data}, null_count_);

        capacity_ = length_ = null_count_ = 0;
        return turbo::OkStatus();
    }

    const uint8_t *FixedSizeBinaryBuilder::get_value(int64_t i) const {
        const uint8_t *data_ptr = byte_builder_.data();
        return data_ptr + i * byte_width_;
    }

    std::string_view FixedSizeBinaryBuilder::get_view(int64_t i) const {
        const uint8_t *data_ptr = byte_builder_.data();
        return {reinterpret_cast<const char *>(data_ptr + i * byte_width_),
                static_cast<size_t>(byte_width_)};
    }

    // ----------------------------------------------------------------------
    // ChunkedArray builders

    namespace internal {

        ChunkedBinaryBuilder::ChunkedBinaryBuilder(int32_t max_chunk_value_length,
                                                   MemoryPool *pool)
                : max_chunk_value_length_(max_chunk_value_length), builder_(new BinaryBuilder(pool)) {
                    DKCHECK_LE(max_chunk_value_length, kBinaryMemoryLimit);
        }

        ChunkedBinaryBuilder::ChunkedBinaryBuilder(int32_t max_chunk_value_length,
                                                   int32_t max_chunk_length, MemoryPool *pool)
                : ChunkedBinaryBuilder(max_chunk_value_length, pool) {
            max_chunk_length_ = max_chunk_length;
        }

        turbo::Status ChunkedBinaryBuilder::finish(ArrayVector *out) {
            if (builder_->length() > 0 || chunks_.size() == 0) {
                std::shared_ptr<Array> chunk;
                TURBO_RETURN_NOT_OK(builder_->finish(&chunk));
                chunks_.emplace_back(std::move(chunk));
            }
            *out = std::move(chunks_);
            return turbo::OkStatus();
        }

        turbo::Status ChunkedBinaryBuilder::NextChunk() {
            std::shared_ptr<Array> chunk;
            TURBO_RETURN_NOT_OK(builder_->finish(&chunk));
            chunks_.emplace_back(std::move(chunk));

            if (auto capacity = extra_capacity_) {
                extra_capacity_ = 0;
                return Reserve(capacity);
            }

            return turbo::OkStatus();
        }

        turbo::Status ChunkedStringBuilder::finish(ArrayVector *out) {
            TURBO_RETURN_NOT_OK(ChunkedBinaryBuilder::finish(out));

            // Change data type to string/utf8
            for (auto &chunk: *out) {
                std::shared_ptr<ArrayData> data = chunk->data()->copy();
                data->type = ::nebula::utf8();
                chunk = std::make_shared<StringArray>(std::move(data));
            }
            return turbo::OkStatus();
        }

        turbo::Status ChunkedBinaryBuilder::Reserve(int64_t values) {
            if (TURBO_UNLIKELY(extra_capacity_ != 0)) {
                extra_capacity_ += values;
                return turbo::OkStatus();
            }

            auto current_capacity = builder_->capacity();
            auto min_capacity = builder_->length() + values;
            if (current_capacity >= min_capacity) {
                return turbo::OkStatus();
            }

            auto new_capacity = BufferBuilder::GrowByFactor(current_capacity, min_capacity);
            if (TURBO_LIKELY(new_capacity <= max_chunk_length_)) {
                return builder_->resize(new_capacity);
            }

            extra_capacity_ = new_capacity - max_chunk_length_;
            return builder_->resize(max_chunk_length_);
        }

    }  // namespace internal

}  // namespace nebula
